#include "DXGICapture.h"
#include "Utils.h"
#include <iostream>

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "dxgi.lib")

DXGICapture::DXGICapture() :
    m_initialized(false),
    m_d3d11Device(nullptr),
    m_d3d11Context(nullptr),
    m_dxgiDuplication(nullptr),
    m_dxgiOutput(nullptr),
    m_dxgiAdapter(nullptr),
    m_dxgiFactory(nullptr) {
    ZeroMemory(&m_textureDesc, sizeof(m_textureDesc));
}

DXGICapture::~DXGICapture() {
    Cleanup();
}

bool DXGICapture::IsSupported() const {
    return Utils::IsWindows8OrLater();
}

bool DXGICapture::Initialize() {
    Utils::DebugLog("DXGICapture: Initializing...");

    if (!IsSupported()) {
        m_lastError = "DXGI capture requires Windows 8 or later";
        return false;
    }

    if (!InitializeD3D11()) {
        return false;
    }

    if (!InitializeDXGI()) {
        Cleanup();
        return false;
    }

    m_initialized = true;
    Utils::DebugLog("DXGICapture: Initialized successfully");
    return true;
}

void DXGICapture::Cleanup() {
    Utils::DebugLog("DXGICapture: Cleaning up...");

    if (m_dxgiDuplication) {
        m_dxgiDuplication->Release();
        m_dxgiDuplication = nullptr;
    }

    if (m_dxgiOutput) {
        m_dxgiOutput->Release();
        m_dxgiOutput = nullptr;
    }

    if (m_dxgiAdapter) {
        m_dxgiAdapter->Release();
        m_dxgiAdapter = nullptr;
    }

    if (m_dxgiFactory) {
        m_dxgiFactory->Release();
        m_dxgiFactory = nullptr;
    }

    if (m_d3d11Context) {
        m_d3d11Context->Release();
        m_d3d11Context = nullptr;
    }

    if (m_d3d11Device) {
        m_d3d11Device->Release();
        m_d3d11Device = nullptr;
    }

    m_initialized = false;
}

bool DXGICapture::InitializeD3D11() {
    D3D_FEATURE_LEVEL featureLevel;
    D3D_FEATURE_LEVEL featureLevels[] = {
        D3D_FEATURE_LEVEL_11_0,
        D3D_FEATURE_LEVEL_10_1,
        D3D_FEATURE_LEVEL_10_0,
        D3D_FEATURE_LEVEL_9_3
    };

    HRESULT hr = D3D11CreateDevice(
        nullptr,
        D3D_DRIVER_TYPE_HARDWARE,
        nullptr,
        0,
        featureLevels,
        ARRAYSIZE(featureLevels),
        D3D11_SDK_VERSION,
        &m_d3d11Device,
        &featureLevel,
        &m_d3d11Context
    );

    if (FAILED(hr)) {
        m_lastError = "Failed to create D3D11 device";
        return false;
    }

    return true;
}

bool DXGICapture::InitializeDXGI() {
    HRESULT hr = CreateDXGIFactory1(__uuidof(IDXGIFactory1), (void**)&m_dxgiFactory);
    if (FAILED(hr)) {
        m_lastError = "Failed to create DXGI factory";
        return false;
    }

    hr = m_dxgiFactory->EnumAdapters1(0, &m_dxgiAdapter);
    if (FAILED(hr)) {
        m_lastError = "Failed to enumerate DXGI adapter";
        return false;
    }

    hr = m_dxgiAdapter->EnumOutputs(0, (IDXGIOutput**)&m_dxgiOutput);
    if (FAILED(hr)) {
        m_lastError = "Failed to enumerate DXGI output";
        return false;
    }

    hr = m_dxgiOutput->DuplicateOutput(m_d3d11Device, &m_dxgiDuplication);
    if (FAILED(hr)) {
        if (hr == DXGI_ERROR_NOT_CURRENTLY_AVAILABLE) {
            m_lastError = "Desktop duplication not available (probably already in use)";
        } else {
            m_lastError = "Failed to create desktop duplication";
        }
        return false;
    }

    DXGI_OUTDUPL_DESC duplicateDesc;
    m_dxgiDuplication->GetDesc(&duplicateDesc);

    m_textureDesc.Width = duplicateDesc.ModeDesc.Width;
    m_textureDesc.Height = duplicateDesc.ModeDesc.Height;
    m_textureDesc.MipLevels = 1;
    m_textureDesc.ArraySize = 1;
    m_textureDesc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
    m_textureDesc.SampleDesc.Count = 1;
    m_textureDesc.SampleDesc.Quality = 0;
    m_textureDesc.Usage = D3D11_USAGE_STAGING;
    m_textureDesc.BindFlags = 0;
    m_textureDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    m_textureDesc.MiscFlags = 0;

    return true;
}

bool DXGICapture::CaptureScreen(HBITMAP& hBitmap, int& width, int& height) {
    if (!m_initialized) {
        m_lastError = "DXGI capture not initialized";
        return false;
    }

    IDXGIResource* desktopResource = nullptr;
    DXGI_OUTDUPL_FRAME_INFO frameInfo;

    HRESULT hr = m_dxgiDuplication->AcquireNextFrame(1000, &frameInfo, &desktopResource);
    if (FAILED(hr)) {
        if (hr == DXGI_ERROR_WAIT_TIMEOUT) {
            m_lastError = "No new frame available";
        } else {
            m_lastError = "Failed to acquire next frame";
        }
        return false;
    }

    ID3D11Texture2D* acquiredDesktopImage = nullptr;
    hr = desktopResource->QueryInterface(__uuidof(ID3D11Texture2D), (void**)&acquiredDesktopImage);
    if (FAILED(hr)) {
        m_lastError = "Failed to query desktop texture interface";
        desktopResource->Release();
        m_dxgiDuplication->ReleaseFrame();
        return false;
    }

    ID3D11Texture2D* stagingTexture = nullptr;
    hr = m_d3d11Device->CreateTexture2D(&m_textureDesc, nullptr, &stagingTexture);
    if (FAILED(hr)) {
        m_lastError = "Failed to create staging texture";
        acquiredDesktopImage->Release();
        desktopResource->Release();
        m_dxgiDuplication->ReleaseFrame();
        return false;
    }

    m_d3d11Context->CopyResource(stagingTexture, acquiredDesktopImage);

    hBitmap = ConvertToHBitmap(stagingTexture);
    width = m_textureDesc.Width;
    height = m_textureDesc.Height;

    stagingTexture->Release();
    acquiredDesktopImage->Release();
    desktopResource->Release();
    m_dxgiDuplication->ReleaseFrame();

    return (hBitmap != nullptr);
}

HBITMAP DXGICapture::ConvertToHBitmap(ID3D11Texture2D* texture) {
    D3D11_MAPPED_SUBRESOURCE mappedResource;
    HRESULT hr = m_d3d11Context->Map(texture, 0, D3D11_MAP_READ, 0, &mappedResource);
    if (FAILED(hr)) {
        m_lastError = "Failed to map staging texture";
        return nullptr;
    }

    BITMAPINFO bmi = { 0 };
    bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmi.bmiHeader.biWidth = m_textureDesc.Width;
    bmi.bmiHeader.biHeight = -static_cast<LONG>(m_textureDesc.Height);
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biBitCount = 32;
    bmi.bmiHeader.biCompression = BI_RGB;

    HDC hdc = GetDC(nullptr);
    if (!hdc) {
        m_d3d11Context->Unmap(texture, 0);
        m_lastError = "Failed to get device context";
        return nullptr;
    }

    void* dibBits = nullptr;
    HBITMAP hBitmap = CreateDIBSection(hdc, &bmi, DIB_RGB_COLORS, &dibBits, nullptr, 0);

    if (hBitmap && dibBits) {
        const BYTE* srcData = static_cast<const BYTE*>(mappedResource.pData);
        BYTE* dstData = static_cast<BYTE*>(dibBits);

        for (UINT y = 0; y < m_textureDesc.Height; ++y) {
            memcpy(dstData + y * m_textureDesc.Width * 4,
                   srcData + y * mappedResource.RowPitch,
                   m_textureDesc.Width * 4);
        }
    } else {
        m_lastError = "Failed to create DIB section";
    }

    ReleaseDC(nullptr, hdc);
    m_d3d11Context->Unmap(texture, 0);

    return hBitmap;
}